#!/usr/bin/env python

# example treemodelsort.py

import os
import sys

import pygtk
pygtk.require('2.0')
import gtk
import gobject

import hashlib
import shutil
import tarfile

import subprocess

NoneHash = '0' * 32

g_debug = False

def md5_file(f, block_size=2**20):
    md5 = hashlib.md5()
    while True:
        data = f.read(block_size)
        if not data:
            break
        md5.update(data)
    return md5.hexdigest()

def tar_folder(folder, target):
    tf = tarfile.open(target, 'w:bz2')
    tf.add(folder)
    tf.close()

class ProjectContent:
    def __init__(self, path):
        self.__files = {}
        self.__path = path
        self.__index_file = os.path.join(path, 'bf.index')
        self.__blobs_dir = os.path.join(path, 'bf.blobs')

    def get_index_file(self):
        return self.__index_file

    def get_blobs_dir(self):
        return self.__blobs_dir

    def get_path(self):
        return self.__path

    def get_files(self):
        return self.__files

    def add_path(self, path, hash = NoneHash):
        if path != '':
            self.__files[path.strip()] = hash.strip()

    def del_path(self, path):
        del self.__files[path]

class MainWindow(gtk.Window):
    __project_content = None

    # changed against the original project
    def load(self):
        files = {}

        try:
            f = open(self.__project_content.get_index_file(), 'a+')
        except e:
            print e

        lines = f.readlines()
        for line in lines:
            ''' line = hash + ' ' + filename '''
            if len(line) < 33:
                continue

            hash = line[:32]
            path = line[33:]

            if path != '':
                files[path.strip()] = hash.strip()
        f.close()
        return files

    def save(self):
        try:
            f = open(self.__project_content.get_index_file(), 'w+')
        except e:
            return

        for k, v in self.__project_content.get_files().items():
            line = '%s %s\n' % (v, k)
            f.writelines([line])

        f.close()

    def __on_quit(self, win, evt = gtk.gdk.DELETE):
        gtk.main_quit(win)

    def __change_title(self, title='Restore the shit'):
        ''' Change the window title to current project path '''
        self.set_title(title)

    def __init__(self):
        gtk.Window.__init__(self)
        self.connect('delete_event', self.__on_quit)

        self.__change_title()

        self.set_default_size(800, 480)
        self.set_size_request(800, 480)

        self.__body_vbox = gtk.VBox()
        self.add(self.__body_vbox)

        self.__create_proj_list()

        settings = self.get_settings()
        settings.set_property("gtk-button-images", True)

        self.__open_project(os.path.realpath(os.path.dirname(__file__)))

        self.show_all()

    def __load_files_to_view(self):
        files = self.load()
        self.__add_files(files)

    def __open_project(self, path):
        ''' Open or create a project_content, should close previous first '''

        self.__proj_path = path

        self.__project_content = ProjectContent(path)
        self.__load_files_to_view()

    def __backup_project_to_file(self, target_fn):
        if not target_fn.endswith('.tar.bz2'):
            target_fn = target_fn + '.tar.bz2'

        blobdir = self.__project_content.get_blobs_dir()

        # empty the target directory
        if os.path.exists(blobdir):
            if os.path.isdir(blobdir):
                shutil.rmtree(blobdir)
            else:
                os.remove(blobdir)

        os.mkdir(blobdir)

        for path, v in self.__project_content.get_files().items():
            try:
                fh = open(path, 'r')
            except:
                continue
            hash = md5_file(fh)
            fh.close()

            self.__project_content.add_path(path, hash)

            target_blob = os.path.join(blobdir, hash)
            if not os.path.exists(target_blob):
                shutil.copy(path, target_blob)

        self.save()

        curdir = os.curdir
        os.chdir(os.path.dirname(self.__proj_path))
        tar_folder(os.path.basename(self.__proj_path), target_fn)
        os.chdir(curdir)

    def __on_backup_project(self, button):
        if not self.__project_content:
            return

        file_dialog = gtk.FileChooserDialog(
            title = 'Select the target zip file',
            action = gtk.FILE_CHOOSER_ACTION_SAVE,
            buttons = (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))

        if file_dialog.run() == gtk.RESPONSE_OK:
            target_fn = file_dialog.get_filename()
            file_dialog.destroy()

            self.__backup_project_to_file(target_fn)

            dlg = gtk.MessageDialog(None, gtk.DIALOG_MODAL, gtk.MESSAGE_WARNING, gtk.BUTTONS_OK)
            dlg.set_markup('<big><b>%s</b></big>' % 'Backup complete')
            dlg.format_secondary_markup('From [%s] to [%s].' % (self.__proj_path, target_fn))
            dlg.run()
            dlg.destroy()

        else:
            file_dialog.destroy()

    def __add_files(self, files):
        model = self.__sort_model.get_model()
        cur_files = self.__project_content.get_files()

        for f, h in files.items():
            if os.path.isdir(f):
                size = -1

                if not f.endswith('/'):
                    f = f + '/'
            else:
                try:
                    size = os.path.getsize(f)
                except:
                    size = -1

            if cur_files.has_key(f):
                continue

            col = [ '', f, os.path.splitext(f)[1], size]
            model.append(col)
            self.__project_content.add_path(f, h)

            # add the parent directory also
            if os.path.isdir(f):
                f = f[:-1]
            dirname = os.path.dirname(f)
            if dirname != '':
                self.__add_files({dirname : NoneHash})

    def add_directory(self, directory):
        self.__add_files({directory : NoneHash})

        try:
            ls = os.listdir(directory)
        except:
            if g_debug:
                print 'access deny'
        else:
            for l in ls:
                temp = os.path.join(directory, l)
                if (os.path.isdir(temp)):
                    self.add_directory(temp)
                else:
                    self.__add_files({temp : NoneHash})

    def del_selected(self):

        def match_text(model, path, iter, data):
            [selection, text] = data
            path = model.get(iter, 1)[0]

            # match dir
            if text.endswith('/'):
                if path.startswith(text):
                    selection.select_iter(iter)
            elif path == text:
                    selection.select_iter(iter)

        selection = self.__proj_tree_view.get_selection()
        sorted_model, selected = selection.get_selected_rows()
        model = sorted_model.get_model()
        iters = [sorted_model.convert_iter_to_child_iter(None, sorted_model.get_iter(path)) for path in selected]

        selection.unselect_all()

        for iter in iters:
            path = model.get(iter, 1)[0]
            self.__proj_tree_view.get_model().foreach(match_text, [selection, path])

        sorted_model, selected = selection.get_selected_rows()
        model = sorted_model.get_model()
        iters = [sorted_model.convert_iter_to_child_iter(None, sorted_model.get_iter(path)) for path in selected]
        for iter in iters:
            path = model.get(iter, 1)[0]
            self.__project_content.del_path(path)
            model.remove(iter)

    def __on_tree_view_key_press(self, treeview, key):
        if key.keyval == gtk.keysyms.space:
            selection = self.__proj_tree_view.get_selection()
            sorted_model, selected = selection.get_selected_rows()
            model = sorted_model.get_model()
            iters = [sorted_model.convert_iter_to_child_iter(None, sorted_model.get_iter(path)) for path in selected]

            blobdir = self.__project_content.get_blobs_dir()
            files = self.__project_content.get_files()
            for iter in iters:
                path = model.get(iter, 1)[0]
                if not path.endswith('/'):
                    path = os.path.join(blobdir, files[path])
                    if os.path.exists(path):
                        subprocess.Popen(['gnome-open', path])

    def __mkdir(self, basedir, newdir):
        for d in newdir.split('/'):
            if d:
                basedir = os.path.join(basedir, d)
                try:
                    os.mkdir(basedir)
                except:
                    pass


    def __copy(self, blob, dest):
        if g_debug:
            print '__copy:: %s --> %s' % (blob, dest)
        shutil.copy(blob, dest)

    def __restore(self, path):
        dialog = RestoreWindow(path)
        if dialog.run() == gtk.RESPONSE_OK:
            ''' src_full = path = "/home/auv/shit" '''
            src_full = path

            ''' src_select = "/home/auv", not include end / '''
            src_select = dialog.get_source_path()[:-1]

            ''' src_drop = "/home" '''
            src_drop = src_select[:src_select.rfind('/')]

            ''' src_keep = "auv/shit" '''
            src_keep = src_full[src_select.rfind('/')+1:]

            ''' dst = "/usr", not include end / '''
            dst_root = dialog.get_target_path()

            blobdir = self.__project_content.get_blobs_dir()
            files = self.__project_content.get_files()

            if src_keep == '/':
                src_keep = ''
            if src_full == '/':
                src_full = ''

            for f, h in files.items():
                if g_debug:
                    print 'src_full:: ' + src_full
                    print "for:: " + f
                if not f.startswith(src_full):
                    continue

                if src_full.endswith('/'):
                    src_full = src_full[:-1]

                if f != src_full and f[len(src_full)] != '/':
                    continue

                blob = os.path.join(blobdir, h)
                new_path = f[len(src_drop) + 1:]
                dst_full = os.path.join(dst_root, new_path)

                if blob.endswith(NoneHash):
                    self.__mkdir(dst_root, new_path)
                else:
                    slashpos = new_path.rfind('/')
                    if slashpos >= 0:
                        self.__mkdir(dst_root, new_path[:slashpos])
                    self.__copy(blob, dst_full)

        dialog.destroy()

    def __on_tree_view_button_press(self, treeview, event):
        if event.button == 1:
            if event.type == gtk.gdk._2BUTTON_PRESS:
                pathinf = treeview.get_path_at_pos(int(event.x), int(event.y))
                if pathinf:
                    path, col, x, y = pathinf
                    model = treeview.get_model()
                    fn = model.get_value(model.get_iter(path), 1)

                    self.__restore(fn)
            else:
                selection = self.__proj_tree_view.get_selection()
                selection.unselect_all()

    def __create_proj_list(self):
        scroll_window = gtk.ScrolledWindow()
        scroll_window.set_border_width(4)
        scroll_window.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
        scroll_window.set_shadow_type(gtk.SHADOW_NONE)

        self.__body_vbox.pack_start(scroll_window)

        self.__proj_list_store = gtk.ListStore(str, str, str, int)

        self.__sort_model = gtk.TreeModelSort(self.__proj_list_store)
        self.__sort_model.set_sort_column_id(1, gtk.SORT_ASCENDING)

        self.__proj_tree_view = gtk.TreeView(self.__sort_model)
        scroll_window.add(self.__proj_tree_view)

        self.__proj_tree_view.set_rules_hint(True)

        self.__proj_tree_view.get_selection().set_mode(gtk.SELECTION_SINGLE)
        self.__proj_tree_view.connect('key-press-event', self.__on_tree_view_key_press)
        self.__proj_tree_view.connect('button-press-event', self.__on_tree_view_button_press)

        self.__proj_tree_view.column = [None]*4
        self.__proj_tree_view.column[0] = gtk.TreeViewColumn('S')
        self.__proj_tree_view.column[1] = gtk.TreeViewColumn('Path')
        self.__proj_tree_view.column[1].set_resizable(True)
        self.__proj_tree_view.column[1].set_expand(True)
        self.__proj_tree_view.column[2] = gtk.TreeViewColumn('Ext')
        self.__proj_tree_view.column[2].set_resizable(True)
        self.__proj_tree_view.column[3] = gtk.TreeViewColumn('Size')
        self.__proj_tree_view.column[3].set_resizable(True)
        self.__proj_tree_view.cell = [None]*4
        for i in range(4):
            self.__proj_tree_view.cell[i] = gtk.CellRendererText()
            self.__proj_tree_view.append_column(self.__proj_tree_view.column[i])
            self.__proj_tree_view.column[i].set_sort_column_id(i)
            self.__proj_tree_view.column[i].pack_start(self.__proj_tree_view.cell[i], True)
            self.__proj_tree_view.column[i].set_attributes(self.__proj_tree_view.cell[i], text=i)

class RestoreWindow(gtk.Dialog):
    def __on_select_node(self, button):
        if not button.get_active():
            return ''

        full_path = ''
        for b in self.__path_nodes[1:]:
            text = b.get_label()
            full_path += '/' + text

            if b == button:
                break

        if button == self.__path_nodes[0]:
            full_path = '/'

        for b in self.__path_nodes:
            if b == button:
                continue
            b.set_active(False)

        ol = len(self.__path)
        cl = len(full_path)
        if ol != cl and self.__path[cl] == '/':
            full_path += '/'

        self.__target_path = full_path
        if g_debug:
            print full_path

    def get_source_path(self):
        return self.__target_path

    def get_target_path(self):
        return self.chooser.get_filename()

    def __init__(self, path):
        gtk.Dialog.__init__(self)

        self.__path = path
        self.__path_nodes = []
        self.__target_path = ''

        frame = gtk.Frame('Restore from ...')
        self.vbox.pack_start(frame, False, True, 0)

        self.hbox = gtk.HBox()
        frame.add(self.hbox)

        path_node = gtk.ToggleButton('/')
        self.hbox.pack_start(path_node, expand=False)
        self.__path_nodes.append(path_node)
        path_node.connect('toggled', self.__on_select_node)

        for dot in path[1:].split('/'):
            if dot:
                path_node = gtk.ToggleButton(dot)
                self.hbox.pack_start(path_node, expand=False)
                self.__path_nodes.append(path_node)
                path_node.connect('toggled', self.__on_select_node)
        self.__path_nodes[-1].set_active(True)

        self.vbox.pack_start(gtk.HSeparator(), expand=False, padding=8)

        frame = gtk.Frame('To ...')
        self.vbox.pack_start(frame)

        self.chooser = gtk.FileChooserWidget(gtk.FILE_CHOOSER_ACTION_SELECT_FOLDER)

        frame.add(self.chooser)

        self.add_buttons(gtk.STOCK_SAVE, gtk.RESPONSE_OK, gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL)

        self.show_all()

if __name__ == '__main__':
    if '--debug' in sys.argv:
        g_debug = True

    gobject.threads_init()
    MainWindow()
    gtk.main()
